package edu.vub.at.kryo;

import java.io.ObjectStreamException;
import java.io.PrintWriter;
import java.io.StringWriter;
import java.lang.reflect.Constructor;
import java.lang.reflect.Field;
import java.lang.reflect.InvocationHandler;
import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Proxy;
import java.util.Set;
import java.util.UUID;

import com.esotericsoftware.kryo.Kryo;
import com.esotericsoftware.kryo.Registration;
import com.esotericsoftware.kryo.Serializer;
import com.esotericsoftware.kryo.io.Input;
import com.esotericsoftware.kryo.io.Output;
import com.esotericsoftware.minlog.Log;
import com.esotericsoftware.minlog.Log.Logger;

import edu.vub.at.actors.natives.SerializationTest.TEST;
import edu.vub.at.eval.Evaluator;
import edu.vub.at.eval.PartialBinder;
import edu.vub.at.exceptions.InterpreterException;
import edu.vub.at.objects.ATObject;
import edu.vub.at.objects.ATTable;
import edu.vub.at.objects.coercion.Coercer;
import edu.vub.at.objects.grammar.ATSymbol;
import edu.vub.at.objects.mirrors.NativeClosure;
import edu.vub.at.objects.mirrors.PrimitiveMethod;
import edu.vub.at.objects.natives.NATBoolean;
import edu.vub.at.objects.natives.NATFraction;
import edu.vub.at.objects.natives.NATNil;
import edu.vub.at.objects.natives.NATNumber;
import edu.vub.at.objects.natives.NATTable;
import edu.vub.at.objects.natives.NATText;
import edu.vub.at.objects.natives.NATTypeTag;
import edu.vub.at.objects.natives.NATTypeTag.OBJRootType;
import edu.vub.at.objects.natives.NativeATObject;
import edu.vub.at.objects.natives.OBJLexicalRoot;
import edu.vub.at.objects.natives.grammar.AGAssignmentSymbol;
import edu.vub.at.objects.natives.grammar.AGSymbol;
import edu.vub.util.MultiMap;

public class ATKryo {
	private Kryo kryo;
	
	static public class MyLogger extends Logger {
		private int IGNORE_LEVEL;
		public MyLogger(int ignore){
			IGNORE_LEVEL= ignore;
		}
		public MyLogger() {
		}
		public void log(int level, String category, String message, Throwable ex) {
			if(level == IGNORE_LEVEL){
				return;
			}
			StringBuilder builder = new StringBuilder(256);
			switch (level) {
            case Log.LEVEL_ERROR:
                    builder.append(" ERROR: ");
                    break;
            case Log.LEVEL_WARN:
                    builder.append("  WARN: ");
                    break;
            case Log.LEVEL_INFO:
                    builder.append("  INFO: ");
                    break;
            case Log.LEVEL_DEBUG:
                    builder.append(" DEBUG: ");
                    break;
            case Log.LEVEL_TRACE:
                    builder.append(" TRACE: ");
                    break;
            }


			builder.append('[');
			builder.append(category);
			builder.append("] ");
			builder.append('[');
			builder.append(Thread.currentThread().getName());
			builder.append("] ");
			builder.append(message);
			if (ex != null) {
				StringWriter writer = new StringWriter(256);
				ex.printStackTrace(new PrintWriter(writer));
				builder.append('\n');
				builder.append(writer.toString().trim());
			}
			System.out.println(builder);
		}
	}

	public ATKryo() throws ClassNotFoundException{
		// new ExtendedMapReferenceResolver()
		kryo = new Kryo(){
			@Override
			protected Serializer newDefaultSerializer(Class type) {
				// Some serialisation logic for anonymous inner classes
				if (type.getSuperclass() == PrimitiveMethod.class) {
					return new PrimitiveMethodSerializer();
				} else if (type.getSuperclass() == PartialBinder.class) {
					return new PartialBinderSerializer();
				} else {
					return newSerializer(ExtendedFieldSerializer.class, type);
					//return newSerializer(FieldSerializer.class, type);
					//return newSerializer(JavaSerializer.class, type);
				}
			}
			
			
		};
		
//		Log.LEVEL_WARN
		//Log.setLogger(new MyLogger(Log.LEVEL_WARN));


		Log.NONE();
		//Log.set(Log.LEVEL_DEBUG);		
		
		
		//Default Serializers
		kryo.addDefaultSerializer(Object[].class, ExtendedObjectSerializers.ObjectArraySerializer.class);
		//kryo.addDefaultSerializer(Collection.class, ExtendedObjectSerializers.CollectionSerializer.class);
		//kryo.addDefaultSerializer(Map.class, ExtendedObjectSerializers.MapSerializer.class);
		
		
		//AG serializers
		kryo.register(AGSymbol.class, new AGSymbolSerializer());
		kryo.register(AGAssignmentSymbol.class, new AGAssignmentSymbolSerializer());
		
		
		//Primitive Serializers
		kryo.register(NATText.class, new NATTextSerializer());
		kryo.register(NATNumber.class, new NATNumberSerializer());
		kryo.register(NATBoolean.class, new NATBooleanSerializer());
		kryo.register(NATFraction.class, new NATFractionSerializer());
		kryo.register(NATTable.class, new NATTableSerializer());
		
		kryo.register(NATTable.EMPTY.getClass(), new NATTableEMPTYSerializer());

		
		//Collection serializers
		kryo.register(MultiMap.class, new MultiMapSerializer());
		
		//Proxy class serializer
		kryo.register(InvocationHandler.class, new KryoProxySerializer());
		
		//UUID serializer
		kryo.register(UUID.class, new UUIDSerializer());
		

		kryo.setAutoReset(true);
		
		
		
		//kryo.register(MethodDictionary.class, new MethodDictionarySerializer());
		//kryo.register(Coercer.class, new CoercerSerializer());				
		//kryo.register(OBJLexicalRoot.class).setSerializer(new OBJLexicalRootSerializer());
		//kryo.register(FieldMap.class, new FieldMapSerializer());		

		
		//kryo.setDefaultSerializer(JavaSerializer.class);
		
	}
	
	public Kryo getKryo(){
		return kryo;
	}
	
	public static Object writeReplace(Object value) throws ObjectStreamException{
		Object prevObject = null;
		
		while(value != null && value instanceof NativeATObject){
			prevObject = value;
			value = ((NativeATObject) value).writeReplace();
			if(prevObject == value) break;
		}
			
		while (value != null && value instanceof TEST) {
			prevObject = value;
			value = ((TEST) value).writeReplace();
			if(prevObject == value) break;
		}
		
		if(prevObject != null){
			return prevObject;
		} else{
			return value;
		}
	}
	
	public static Object readResolve(Object value) throws ObjectStreamException{
		if (value != null && value instanceof NativeATObject) {
			value = ((NativeATObject) value).readResolve();
		}

		if (value != null && value instanceof TEST) {
			value = ((TEST) value).readResolve();
		}
		return value;
	}

	
	public class UUIDSerializer extends Serializer<UUID> {

	    @Override
	    public void write(Kryo kryo, Output output, UUID u) {
	        output.writeLong(u.getMostSignificantBits());
	        output.writeLong(u.getLeastSignificantBits());
	    }

	    @Override
	    public UUID read(Kryo kryo, Input input, Class<UUID> k) {
	        long msb = input.readLong();
	        long lsb = input.readLong();
	        return new UUID(msb, lsb);
	    }
	}

	
	class NATTableEMPTYSerializer extends Serializer<Object>{

		@Override
		public void write(Kryo kryo, Output output, Object object) {
			Log.warn("Writing NATTableEMPTY with class: " + object.getClass());
			kryo.writeClass(output, object.getClass());
		}

		@Override
		public Object read(Kryo kryo, Input input, Class<Object> type) {
			kryo.readClass(input);
			Log.warn("Reading NATTableEMPTY --> returning NATTable.EMPTY");
			return NATTable.EMPTY;
		}
		
	}
		
	class NativeClosureSerializer extends Serializer<NativeClosure>{

		@Override
		public void write(Kryo kryo, Output output, NativeClosure object) {
			kryo.writeClass(output, object.getClass());
			Log.warn("Writing PrimitiveMethod with class: " + object.getClass());
			Log.warn("Writing PrimitiveMethod with Object: " + NativeClosure.nativeClosureMap.get(object.getClass()));
		}

		@Override
		public NativeClosure read(Kryo kryo, Input input, Class<NativeClosure> type) {
			Registration res = kryo.readClass(input);
			Log.warn("Reading PrimitiveMethod with Class: " + res.getType());
			Log.warn("Reading PrimitiveMethod with Object: " + NativeClosure.nativeClosureMap.get(res.getType()));
			return (NativeClosure) NativeClosure.nativeClosureMap.get(res.getType());
		}

	}
	
	class PrimitiveMethodSerializer extends Serializer<PrimitiveMethod>{

		@Override
		public void write(Kryo kryo, Output output, PrimitiveMethod object) {
			kryo.writeClass(output, object.getClass());
			Log.warn("Writing PrimitiveMethod with class: " + object.getClass());
			Log.warn("Writing PrimitiveMethod with Object: " + PrimitiveMethod.primitiveMap.get(object.getClass()));
		}

		@Override
		public PrimitiveMethod read(Kryo kryo, Input input, Class<PrimitiveMethod> type) {
			Registration res = kryo.readClass(input);
			Log.warn("Reading PrimitiveMethod with Class: " + res.getType());
			Log.warn("Reading PrimitiveMethod with Object: " + PrimitiveMethod.primitiveMap.get(res.getType()));
			return (PrimitiveMethod) PrimitiveMethod.primitiveMap.get(res.getType());
		}

	}
	
	class PartialBinderSerializer extends Serializer<Object>{

		@Override
		public void write(Kryo kryo, Output output, Object object) {
			kryo.writeClass(output, object.getClass());
			Log.warn("Writing PartialBinder with class: " + object.getClass());
			Log.warn("Writing PartialBinder with Object: " + PartialBinder.partialBinderMap.get(object.getClass()));
		}

		@Override
		public Object read(Kryo kryo, Input input, Class<Object> type) {
			Registration res = kryo.readClass(input);

			Log.warn("Reading PartialBinder with Class: " + res);
			Log.warn("Reading PartialBinder with Object: " + PartialBinder.partialBinderMap.get(res.getType()));
			return PartialBinder.partialBinderMap.get(res.getType());
		}

	}
	
	
	class OBJRootTypeSerializer extends Serializer<OBJRootType>{

		@Override
		public void write(Kryo kryo, Output output, OBJRootType object) {
			kryo.writeClass(output, object.getClass());
		}

		@Override
		public OBJRootType read(Kryo kryo, Input input, Class<OBJRootType> type) {
			kryo.readClass(input);
			return OBJRootType._INSTANCE_;
		}
	}


	class NATTypeTagSerializer extends Serializer<NATTypeTag>{

		@Override
		public void write(Kryo kryo, Output output, NATTypeTag object) {
			try {
				kryo.writeClassAndObject(output, object.base_typeName());
				kryo.writeClassAndObject(output, object.base_superTypes());
			} catch (InterpreterException e) {
				e.printStackTrace();
			}
		}

		@Override
		public NATTypeTag read(Kryo kryo, Input input,
				Class<NATTypeTag> type) {		
			ATSymbol typename = (ATSymbol) kryo.readClassAndObject(input);
			ATTable parentTypes = (ATTable) kryo.readClassAndObject(input);
			return NATTypeTag.atValue(typename, parentTypes);
		}
		
	}
	
	
			
	class OBJLexicalRootSerializer extends Serializer<OBJLexicalRoot>{

		@Override
		public void write(Kryo kryo, Output output, OBJLexicalRoot object) {}

		@Override
		public OBJLexicalRoot read(Kryo kryo, Input input,
				Class<OBJLexicalRoot> type) {
			return OBJLexicalRoot._INSTANCE_;
		}
		
	}
	

	
	class NATNilSerializer extends Serializer<NATNil>{

		@Override
		public void write(Kryo kryo, Output output, NATNil object) {}

		@Override
		public NATNil read(Kryo kryo, Input input, Class<NATNil> type) {
			return Evaluator.getNil();
		}
	}

	
	class NATTableSerializer extends Serializer<NATTable>{

		@Override
		public void write(Kryo kryo, Output output, NATTable natTable) {
			kryo.writeClassAndObject(output, natTable.elements_	);
		}

		@Override
		public NATTable read(Kryo kryo, Input input, Class<NATTable> type) {
			return NATTable.atValue((ATObject[])kryo.readClassAndObject(input));
		}
		
	}
	
	
	class CoercerSerializer extends Serializer<Coercer>{

		@Override
		public void write(Kryo kryo, Output output, Coercer object) {
			Log.warn("B");
			Field f;
			try {
				f = object.getClass().getDeclaredField("principal_");

				f.setAccessible(true);

				Log.warn(f.get(object).toString());
				
				Object principal_ = f.get(object); // IllegalAccessException
				
				Log.warn("Writing principal_ with Class: " + principal_.getClass());
				Log.warn("Writing principal_ with Object: " + principal_);
				
				kryo.writeClass(output, principal_.getClass());
				kryo.writeObject(output, principal_);
				
			} catch (NoSuchFieldException | SecurityException e) {
				e.printStackTrace();
			} catch (IllegalArgumentException e) {
				e.printStackTrace();
			} catch (IllegalAccessException e) {
				e.printStackTrace();
			}
		}

		@Override
		public Coercer read(Kryo kryo, Input input, Class<Coercer> type) {
			try {
			
			Class clazz = kryo.readClass(input).getType();
			Log.warn("Reading  principal_ Class: " + clazz);
			
			Object principal_ = kryo.readObject(input, clazz);
			
			Log.warn("Reading  principal_ Class: " + principal_.getClass());
			Log.warn("Reading principal_ with Object: " + principal_);
			
			Constructor g = Coercer.class.getDeclaredConstructor(new Class[]{ATObject.class, Thread.class});
			g.setAccessible(true);
			
			return (Coercer) g.newInstance(principal_, Thread.currentThread());
				
			} catch (InstantiationException | IllegalAccessException
					| IllegalArgumentException | InvocationTargetException e) {
				e.printStackTrace();
			} catch (NoSuchMethodException e) {
				e.printStackTrace();
			} catch (SecurityException e) {
				e.printStackTrace();
			}
			return null;
			
		}
		
	}
	
//	class KryoProxySerializer extends Serializer<Object> {
//
//		@Override
//		public void write(Kryo kryo, Output output, Object object) {
//			Log.warn("A");
//			//We write the interfaces
//			kryo.writeObject(output, object.getClass().getInterfaces());	
//			//We write the InvocationHandler
//			Log.warn(((InvocationHandler) Proxy.getInvocationHandler(object)).getClass().getName());
//			kryo.writeClassAndObject( output, ((InvocationHandler) Proxy.getInvocationHandler(object)));
//			
//		}
//
//		@Override
//		public Object read(Kryo kryo, Input input, Class<Object> type) {
//			Log.warn("C");
//			//We read the interfaces
//			final Class<?>[] interfaces = kryo.readObject(input, Class[].class);
//			//We read the InvocationHandler
//			final InvocationHandler proxyHandler = (InvocationHandler) kryo.readClassAndObject(input);	
//			//We return a new proxy class
//			return createProxy(kryo, proxyHandler, interfaces);
//		}
//
//		private Object createProxy(Kryo kryo, InvocationHandler proxyHandler, Class<?>[] interfaces) {
//			return Proxy.newProxyInstance(Thread.currentThread().getContextClassLoader(), interfaces, proxyHandler);
//			//return Proxy.newProxyInstance(kryo.getClassLoader(), interfaces, proxyHandler);
//		}
//
//	}
	
	//is called when kryo finds out an object is of type Proxy
	class KryoProxySerializer extends Serializer<Object> {

		@Override
		public void write(Kryo kryo, Output output, Object object) {
			//We write the interfaces
			kryo.writeObject(output, object.getClass().getInterfaces());	
			//We write the InvocationHandler
			Log.warn(((InvocationHandler) Proxy.getInvocationHandler(object)).getClass().getName());
			Coercer coercer = (Coercer) Proxy.getInvocationHandler(object);
			Log.warn(coercer.toString());

			try {
				//We only need to write the principal_ field
				
//				Field f = coercer.getClass().getDeclaredField("principal_");
//				f.setAccessible(true);
//				Object principal_ = f.get(coercer);
				ATObject principal_ = coercer.getPrincipal();
				
				kryo.writeClassAndObject(output, principal_);
				
			} catch (SecurityException e) {
				e.printStackTrace();
			} catch (IllegalArgumentException e) {
				e.printStackTrace();
			}
			
		}
		

		@Override
		public Object read(Kryo kryo, Input input, Class<Object> type) {	
			//We read the interfaces
			final Class<?>[] interfaces = kryo.readObject(input, Class[].class);

			// We read the InvocationHandler
			try {
				Object principal_ = kryo.readClassAndObject(input);
				Constructor g = Coercer.class.getDeclaredConstructor(new Class[]{ATObject.class, Thread.class});
				g.setAccessible(true);
				
				Coercer proxyHandler = (Coercer) g.newInstance(principal_, Thread.currentThread());
				// We return a new proxy class
				return createProxy(kryo, proxyHandler, interfaces);
			} catch (NoSuchMethodException | SecurityException e) {
				e.printStackTrace();
			} catch (InstantiationException e) {
				e.printStackTrace();
			} catch (IllegalAccessException e) {
				e.printStackTrace();
			} catch (IllegalArgumentException e) {
				e.printStackTrace();
			} catch (InvocationTargetException e) {
				e.printStackTrace();
			}
			return null; //Beter een error throwe
		}

		private Object createProxy(Kryo kryo, InvocationHandler proxyHandler, Class<?>[] interfaces) {
			return Proxy.newProxyInstance(Thread.currentThread().getContextClassLoader(), interfaces, proxyHandler);
			//return Proxy.newProxyInstance(kryo.getClassLoader(), interfaces, proxyHandler);
		}

	}
	
	class MultiMapSerializer extends Serializer<MultiMap>{ //TODO kan beter

		@Override
		public void write(Kryo kryo, Output output, MultiMap object) {
			Set values = object.entrySet();
		
			output.writeInt(values.size());
			for(Object element : values){	
				kryo.writeClassAndObject(output, ((MultiMap.Entry) element).getKey());
				kryo.writeClassAndObject(output, ((MultiMap.Entry) element).getValue());
			}
			
		}

		@Override
		public MultiMap read(Kryo kryo, Input input, Class<MultiMap> type) {
			int ctr = input.readInt();
			MultiMap mp = new MultiMap();
			while(ctr > 0){
				mp.put(kryo.readClassAndObject(input), kryo.readClassAndObject(input));
				ctr--;
			}
			return mp;	
		}
		
	}
	
	
//	class FieldMapSerializer extends Serializer<FieldMap>{
//
//		@Override
//		public void write(Kryo kryo, Output output, FieldMap fieldMap) {
//			ATSymbol[] varnames = fieldMap.listFields();
//			output.writeInt(varnames.length);
//			
//			for(int i = 0; i < varnames.length; i++){
//				kryo.writeClassAndObject(output, varnames[i]);
//			}
//		}
//
//		@Override
//		public FieldMap read(Kryo kryo, Input input, Class<FieldMap> type) {
//			FieldMap res = new FieldMap();
//			int length = input.readInt();
//			
//			for(int i = 0; i < length; i++){
//				res.put((ATSymbol) kryo.readClassAndObject(input));
//			}
//			return res;
//		}
//	}
	
	class AGSymbolSerializer extends Serializer<AGSymbol> { //

		@Override
		public void write(Kryo kryo, Output output, AGSymbol object) {
			// since symbols are interned, they are shared among parse trees
			// (even among actors!) so their source location is meaningless
			output.writeString(object.toString());
			
		}

		@Override
		public AGSymbol read(Kryo kryo, Input input, Class<AGSymbol> type) {
			String txt = input.readString();
			return AGSymbol.jAlloc(txt);
		}
	}
	
	class AGAssignmentSymbolSerializer extends Serializer<AGSymbol> { //

		@Override
		public void write(Kryo kryo, Output output, AGSymbol object) {
			// since symbols are interned, they are shared among parse trees
			// (even among actors!) so their source location is meaningless
			output.writeString(object.toString());
			
		}

		@Override
		public AGSymbol read(Kryo kryo, Input input, Class<AGSymbol> type) {
			String txt = input.readString();
			return AGAssignmentSymbol.jAlloc(txt);
		}
	}
	
	
	class NATTextSerializer extends Serializer<NATText> {

		@Override
		public void write(Kryo kryo, Output output, NATText object) {
			output.writeString(object.javaValue);
			
		}

		@Override
		public NATText read(Kryo kryo, Input input, Class<NATText> type) {
			return NATText.atValue(input.readString());
		}
		
	}
	
	class NATNumberSerializer extends Serializer<NATNumber>{

		@Override
		public void write(Kryo kryo, Output output, NATNumber object) {
			output.writeInt(object.javaValue);
			
		}

		@Override
		public NATNumber read(Kryo kryo, Input input, Class<NATNumber> type) {
			return NATNumber.atValue(input.readInt());
		}
	}
	
	class NATBooleanSerializer extends Serializer<NATBoolean>{

		@Override
		public void write(Kryo kryo, Output output, NATBoolean object) {
			output.writeBoolean(object.javaValue);
		}

		@Override
		public NATBoolean read(Kryo kryo, Input input, Class<NATBoolean> type) {
			return (NATBoolean) NATBoolean.atValue(input.readBoolean());
		}
		
	}
	
	class NATFractionSerializer extends Serializer<NATFraction> {

		@Override
		public void write(Kryo kryo, Output output, NATFraction object) {
			output.writeDouble(object.javaValue);
		}

		@Override
		public NATFraction read(Kryo kryo, Input input, Class<NATFraction> type) {
			return NATFraction.atValue(input.readDouble());
		}
		
	}
	
	
//	class MethodDictionarySerializer extends Serializer<MethodDictionary>{
//
//		@Override
//		public void write(Kryo kryo, Output output, MethodDictionary object) {
//			Set values = object.entrySet();
//			
//			output.writeInt(values.size());
//			for(Object element : values){	
//				kryo.writeClassAndObject(output, ((MethodDictionary.Entry) element).getKey());
//				kryo.writeClassAndObject(output, ((MethodDictionary.Entry) element).getValue());
//			}
//		}
//
//		@Override
//		public MethodDictionary read(Kryo kryo, Input input, Class<MethodDictionary> type) {
//			int ctr = input.readInt();
//			MethodDictionary md = new MethodDictionary();
//			while(ctr > 0){
//				md.put(kryo.readClassAndObject(input), kryo.readClassAndObject(input));
//				ctr--;
//			}
//			return md;	
//		}
//
//	}
	
}
